import mysql.connector
from ksycopg2.extras import DictCursor
from ksycopg2.extras import RealDictCursor
from mysql.connector import Error
import os
import ksycopg2
data_source = None
trace_table_name = os.getenv('DB_TRACE_DATATABLE','trace_data')
DATA_DEFAULT_SOURCE = 'mysql'

def create_connection():
    try:
        global data_source
        data_source = os.getenv('DB_DATA_SOURCE', DATA_DEFAULT_SOURCE)
        if data_source == DATA_DEFAULT_SOURCE:
            connection = mysql.connector.connect(
                host=os.getenv('DB_HOST', '10.21.231.40'),  # 使用环境变量，默认值为 '10.21.231.40'
                port=int(os.getenv('DB_PORT', 3306)),  # 使用环境变量，默认值为 3306
                user=os.getenv('DB_USER', 'trace_user'),  # 使用环境变量，默认值为 'agg_user'
                password=os.getenv('DB_PASSWORD', 'trace_user_password'),  # 使用环境变量，默认值为 'agg_user_password'
                database=os.getenv('DB_DATABASE', 'trace_database')  # 使用环境变量，默认值为 'agg_database'
            )
            if connection.is_connected():
                print('Connected to MySQL database')
                return connection

        # 连接人大金仓数据库
        else:
            connection = ksycopg2.connect(
                host=os.getenv('DB_HOST', '10.161.43.163'),  # 使用环境变量，默认值为 '10.21.231.40'
                port=os.getenv('DB_PORT', '54321'),  # 使用环境变量，默认值为 '54321'
                user=os.getenv('DB_USER', 'system'),  # 使用环境变量，默认值为 'trace_user'
                password=os.getenv('DB_PASSWORD', '123456'),  # 使用环境变量，默认值为 'trace_user_password'
                database=os.getenv('DB_DATABASE', 'fed_database')  # 使用环境变量，默认值为 'trace_database'
            )
            if connection:
                print('Connected to KingbaseES database')
                return connection
    except Error as e:
        print(f"Error: {e}")
        return None

# 修改后的表结构，增加 userid 和 username 字段
def create_trace_data_table(connection):
    cursor = connection.cursor()
    if data_source == DATA_DEFAULT_SOURCE:
        cursor.execute(f"""
            CREATE TABLE IF NOT EXISTS {trace_table_name} (
                id INT AUTO_INCREMENT PRIMARY KEY,
                client_id VARCHAR(255) NOT NULL,
                client_ip VARCHAR(255),
                client_port INT,
                client_sign1 TEXT,
                client_sign2 TEXT,
                client_BFV LONGBLOB,
                pk_sign_1 TEXT,
                pk_sign_2 TEXT,
                userid VARCHAR(255),  -- 增加用户 ID
                username VARCHAR(255),  -- 增加用户名
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
    else:
        cursor.execute(f"""
            CREATE TABLE IF NOT EXISTS {trace_table_name} (
                id SERIAL PRIMARY KEY,
                client_id VARCHAR(255) NOT NULL,
                client_ip VARCHAR(255),
                client_port INT,
                client_sign1 TEXT,
                client_sign2 TEXT,
                client_BFV BLOB,
                pk_sign_1 TEXT,
                pk_sign_2 TEXT,
                userid VARCHAR(255),  -- 增加用户 ID
                username VARCHAR(255),  -- 增加用户名
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
    connection.commit()


# 插入或更新数据时增加 userid 和 username 字段
def insert_or_update_trace_data(connection, client_id, client_ip=None, client_port=None, client_sign1=None,
                                client_sign2=None, client_BFV=None, pk_sign_1=None, pk_sign_2=None,
                                userid=None, username=None):
    cursor = connection.cursor()
    cursor.execute(f"""
        SELECT client_ip, client_port, client_sign1, client_sign2, client_BFV, pk_sign_1, pk_sign_2, userid, username 
        FROM {trace_table_name} WHERE client_id = %s
    """, (client_id,))
    record = cursor.fetchone()

    if record:
        client_ip = client_ip if client_ip is not None else record[0]
        client_port = client_port if client_port is not None else record[1]
        client_sign1 = client_sign1 if client_sign1 is not None else record[2]
        client_sign2 = client_sign2 if client_sign2 is not None else record[3]
        client_BFV = client_BFV if client_BFV is not None else record[4]
        pk_sign_1 = pk_sign_1 if pk_sign_1 is not None else record[5]
        pk_sign_2 = pk_sign_2 if pk_sign_2 is not None else record[6]
        userid = userid if userid is not None else record[7]
        username = username if username is not None else record[8]

        cursor.execute(f"""
            UPDATE {trace_table_name}
            SET client_ip = %s, client_port = %s, client_sign1 = %s, client_sign2 = %s, client_BFV = %s, 
                pk_sign_1 = %s, pk_sign_2 = %s, userid = %s, username = %s, timestamp = CURRENT_TIMESTAMP
            WHERE client_id = %s
        """, (client_ip, client_port, client_sign1, client_sign2, client_BFV, pk_sign_1, pk_sign_2, userid, username,
              client_id))
    else:
        cursor.execute(f"""
            INSERT INTO {trace_table_name} (client_id, client_ip, client_port, client_sign1, client_sign2, client_BFV, 
                                            pk_sign_1, pk_sign_2, userid, username)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """, (client_id, client_ip, client_port, client_sign1, client_sign2, client_BFV, pk_sign_1, pk_sign_2, userid,
              username))

    connection.commit()

def client_id_exists(connection, client_id):
    cursor = connection.cursor()
    cursor.execute(f"""
        SELECT COUNT(*) FROM {trace_table_name} WHERE client_id = %s
    """, (client_id,))
    return cursor.fetchone()[0] > 0
# def get_client_id(connection,id):
#     cursor = connection.cursor()
#     cursor.execute(f"""
#         SELECT client_id FROM {trace_table_name} WHERE id = %s
#     """, (id,))
#     return cursor.fetchone()
# def get_client_info(connection,id):
#     cursor = connection.cursor()
#     cursor.execute(f"""
#         SELECT client_id,userid,username FROM {trace_table_name} WHERE id = %s
#     """, (id,))
#     return cursor.fetchall()


def get_client_info(connection, id):
    cursor = connection.cursor()
    # cursor.execute(f"""
    #     SELECT client_id, userid, username FROM {trace_table_name} WHERE id = %s
    # """, (id,))
    cursor.execute(f"""
        SELECT  username FROM {trace_table_name} WHERE id = %s
    """, (id,))
    result = cursor.fetchone()[0]
    if result:
        username = result
        # return f"username: {username}"
        return username
    else:
        return "No client found for the given id"
    # if result:
    #     client_id, userid, username = result
    #     return f"client_id: {client_id}, userid: {userid}, username: {username}"
    # else:
    #     return "No client found for the given id"


def get_basedata_clients_data(connection):
    cursor = connection.cursor()
    cursor.execute(f"""
        SELECT id, client_id FROM {trace_table_name}
    """)
    return cursor.fetchall()

def get_id(connection,client_id):
    cursor = connection.cursor()
    cursor.execute(f"""
        SELECT id FROM {trace_table_name} WHERE client_id = %s
    """, (client_id,))
    result = cursor.fetchone()
    return result[0] if result else None

def set_id(connection):
    cursor = connection.cursor()

    # 获取所有行并为每一行分配行号
    cursor.execute(f"""
        SELECT row_num, client_id FROM (
            SELECT 
                ROW_NUMBER() OVER (ORDER BY id) AS row_num,
                client_id
            FROM {trace_table_name}
        ) AS numbered
    """)

    rows = cursor.fetchall()
    # 遍历所有的行并更新 id 为 row_num
    for row in rows:
        row_num, client_id = row
        cursor.execute(f"""
            UPDATE {trace_table_name}
            SET id = %s
            WHERE client_id = %s
        """, (row_num, client_id))

    connection.commit()  # 提交更新


def get_trace_data_by_client_id(connection, client_id):
    cursor = connection.cursor()
    cursor.execute(f"""
        SELECT id, client_ip, client_port, client_BFV, pk_sign_1, pk_sign_2 FROM {trace_table_name} WHERE client_id = %s
    """, (client_id,))
    return cursor.fetchone()

def count_trace_clients(connection):
    cursor = connection.cursor()
    cursor.execute(f"""
        SELECT COUNT(*) FROM {trace_table_name}
    """)
    return cursor.fetchone()[0]
def get_simpledata_trace_clients_data(connection):
    if data_source == DATA_DEFAULT_SOURCE:
        cursor = connection.cursor(dictionary=True)
        cursor.execute(f"""
                SELECT id, client_id, userid, username, client_ip FROM {trace_table_name}
            """)
        return cursor.fetchall()
    else:
        cursor = connection.cursor(cursor_factory=RealDictCursor)
        cursor.execute(f"""
            SELECT id, client_id, userid, username, client_ip FROM {trace_table_name}
        """)
        return cursor.fetchall()

def get_basedata_trace_clients_data(connection, client_id):
    cursor = connection.cursor()
    cursor.execute(f"""
        SELECT  client_ip FROM {trace_table_name} WHERE client_id = %s
    """, (client_id,))
    return cursor.fetchall()

def get_data_trace_clients_data(connection):
    if data_source == DATA_DEFAULT_SOURCE:
        cursor = connection.cursor(dictionary=True)
        cursor.execute(f"""
              SELECT id, client_sign1, client_sign2, client_BFV, pk_sign_1, pk_sign_2 FROM {trace_table_name}
          """)
        return cursor.fetchall()
    else:
        cursor = connection.cursor(cursor_factory=RealDictCursor)
        cursor.execute(f"""
            SELECT id, client_sign1, client_sign2, client_BFV, pk_sign_1, pk_sign_2 FROM {trace_table_name}
        """)
        return cursor.fetchall()

def get_all_trace_clients_data(connection):
    cursor = connection.cursor(dictionary=True)
    cursor.execute(f"""
        SELECT id, client_id, client_ip, client_port, client_sign1, client_sign2, client_BFV, pk_sign_1, pk_sign_2 FROM {trace_table_name}
    """)
    return cursor.fetchall()

def count_clients(connection):
    cursor = connection.cursor()
    cursor.execute(f"""
        SELECT COUNT(*) FROM {trace_table_name}
    """)
    return cursor.fetchone()[0]
def delete_client_data(client_id):
    try:
        connection = create_connection()
        cursor = connection.cursor()
        cursor.execute(f"""DELETE FROM {trace_table_name} WHERE client_id = %s""", (client_id,))
        connection.commit()
        print(cursor.rowcount,type(cursor.rowcount))
        return cursor.rowcount  # 返回删除的行数
    except Error as e:
        print(f"Database error: {e}")
        return None  # 返回 None 以表示出现错误
    finally:
        if connection:
            cursor.close()
            connection.close()